# TCP middleman: forward a TCP connection between two peers
# and listen / modify traffic

import socket
import select
import threading
import binascii
import logging

import pyamf

# temporary imports
import random, time

from netutil import *
from rtmp import *                    # RTMP spec constants
from bytebuffer import ByteStream     # tamperable pipe abstraction
from str2int import *                 # strings <-> big/little-endian integers
from fo_client import RTMPStream
import fo_client
from persistentdict import PersistentDict

logging.basicConfig(level=logging.DEBUG,
                    format='%(asctime)s.%(msecs)03d %(lineno)d %(levelname)s %(message)s',
                    datefmt='%H:%M:%S')
log = logging.getLogger('fo')

"""class SocketPair:
    
    def __init__(self, in_, out, in_lock, out_lock):
        self.in_ = in_
        self.out = out
        self.in_buf = ''   # stores recently read data (cleared on each write)
        self.in_lock = in_lock
        self.out_lock = out_lock

    def read(self, nbytes):
        '''
        Returns the next nbytes from the input socket.
        '''
        data = ''
        with self.in_lock:
            while len(data) < nbytes:
                data += self.in_.recv(nbytes - len(data))
        self.in_buf += data
        return data

    def read1(self, nbytes):
        '''
        Reads up to nbytes, making only one call to the underlying socket.
        '''
        with self.in_lock:
            data = self.in_.recv(nbytes)
        self.in_buf += data
        return data

    def write(self, data):
        '''
        Writes all of data to the output socket.
        '''
        with self.out_lock:
            self.out.sendall(data)
        self.in_buf = ''

    def forward(self, nbytes):
        '''
        Reads n bytes from input and writes data to output
        '''
        while nbytes > 0:
            with self.in_lock:
                data = self.in_.recv(nbytes)
            with self.out_lock:
                self.out.sendall(data)
            nbytes -= len(data)
        self.in_buf = ''
"""

class MiddleMan:
    "TCP MiddleMan for fiddling with data over a TCP connection"

    def __init__(self, client_bound_tamperer, server_bound_tamperer):
        self.cli_buf = ByteStream()
        self.srv_buf = ByteStream()        
        self.cli_tamperer = threading.Thread(None, client_bound_tamperer, None, (self.cli_buf,))
        self.srv_tamperer = threading.Thread(None, server_bound_tamperer, None, (self.srv_buf,))

    def mediate(self, middle_srv_port, srv_port, srv_ip):
        "Listen on port srv_port for a connection and forward it to srv_addr"
        # Accept incoming client connection
        self.listen_s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.listen_s.bind(('', middle_srv_port))
        self.listen_s.listen(1)
        self.cli_s, self.cli_addr = self.listen_s.accept()
        print "bound to client", self.cli_addr

        # Initiate connection with server
        self.srv_s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.srv_addr = (srv_ip, srv_port)
        self.srv_s.connect(self.srv_addr)
        print "connected to server", self.srv_addr

        # forward data idefinitely
        self.cli_tamperer.start()
        self.srv_tamperer.start()
        self.forwardAndForge()
        #self.forwardWithLocks()
        #self.forward()

    def forward(self):
        "Forward packets between client and server forever"
        srv_buf = ""  # server-bound data
        cli_buf = ""  # client-bound data
        MSS = 4096
        while self.active:
            readable, writable, err = select.select((self.cli_s, self.srv_s),
                                                    (self.cli_s, self.srv_s),
                                                    (), 1)
            if self.cli_s in readable and self.srv_s in writable:
                if len(cli_buf) < MSS:
                    cli_buf += self.cli_s.recv(MSS)
                bytes_sent = self.srv_s.send(cli_buf)
                cli_buf = cli_buf[bytes_sent:]
                print "cli -> srv (%d bytes)" % bytes_sent

            if self.srv_s in readable and self.cli_s in writable:
                if len(srv_buf) < MSS:
                    srv_buf += self.srv_s.recv(MSS)
                bytes_sent = self.cli_s.send(srv_buf)
                srv_buf = srv_buf[bytes_sent:]
                print "srv -> cli (%d bytes)" % bytes_sent

    def forwardAndForge(self):
        '''Forward packets between client and server forever,
        and pass all server-to-client data through '''
        MSS = 4096
        cli_buf = self.cli_buf         # modifiable client-bound data
        srv_buf = self.srv_buf         # modifiable server-bound data
        cli_s = self.cli_s             # local var (for readability)
        srv_s = self.srv_s             # local var (for readability)
        while 1:
            """readtest = ((cli_buf.in_len() < MSS) and srv_s or ()) +\
                       ((srv_buf.in_len() < MSS) and cli_s or ())
            writetest = (cli_buf.out_len() and cli_s or ()) +\
                        (srv_buf.out_len() and srv_s or ())"""

            readable, writable, err = select.select((cli_s, srv_s),
                                                    (cli_s, srv_s), (), 1)

            #log.debug('select: read: %s; write: %s', readable, writable)
            if srv_s in readable: # and cli_buf.in_len() < MSS:
                data = srv_s.recv(MSS)
                cli_buf.put(data)
                #print "srv -> mid (%d bytes)" % len(data)

            if cli_s in writable and cli_buf.out_len():
                bytes_sent = cli_s.send(cli_buf.tryget(MSS))
                #print "mid -> cli (%d bytes)" % bytes_sent

            if cli_s in readable: # and srv_buf.in_len() < MSS:
                data = cli_s.recv(MSS)
                srv_buf.put(data)
                #print "cli -> mid (%d bytes)" % len(data)

            if srv_s in writable and srv_buf.out_len():
                bytes_sent = srv_s.send(srv_buf.tryget(MSS))
                #print "mid -> srv (%d bytes)" % bytes_sent

    """def forwardWithLocks(self):
        '''Forward packets between client and server forever,
        and pass all server-to-client data through '''
        MSS = 4096

        cli_s = self.cli_s             # local var (for readability)
        srv_s = self.srv_s             # local var (for readability)
        cli_lock = threading.Lock()
        srv_lock = threading.Lock()
        client_bound_pair = SocketPair(srv_s, cli_s, srv_lock, cli_lock)
        server_bound_pair = SocketPair(cli_s, srv_s, cli_lock, srv_lock)
        self.cli_tamperer = threading.Thread(None, client_bound_tamperer, None, (client_bound_pair,))
        self.srv_tamperer = threading.Thread(None, server_bound_tamperer, None, (server_bound_pair,))

        while 1:
            readable, writable, err = select.select((cli_s, srv_s),
                                                    (cli_s, srv_s), (), 1)

            if srv_s in readable:
                data = srv_s.recv(MSS)
                cli_buf.put(data)
                #print "srv -> mid (%d bytes)" % len(data)

            if cli_s in writable and cli_buf.out_len():
                bytes_sent = cli_s.send(cli_buf.tryget(MSS))
                #print "mid -> cli (%d bytes)" % bytes_sent

            if cli_s in readable:
                data = cli_s.recv(MSS)
                srv_buf.put(data)
                #print "cli -> mid (%d bytes)" % len(data)

            if srv_s in writable and srv_buf.out_len():
                bytes_sent = srv_s.send(srv_buf.tryget(MSS))
                #print "mid -> srv (%d bytes)" % bytes_sent"""
    

def tamperClientBound(iobuf):
    '''
    Watches data from server to client and inserts or removes data as desired.
    '''

    # forward handshake packets
    iobuf.forward(1)
    iobuf.forward(1536)
    iobuf.forward(1536)

    cs = RTMPStream(iobuf)

    '''dupdata = ""
    dup_count = 1
    time_start = time.time()
    audio_count = 0
    dropped_packets = 0'''

    # Decode next several packets
    for i in range(0):
        packet = rtmp.decode()
        
        # Count num. of AUDIO packets / sec.
        '''if mtype == MTYPE_AUDIO:
            if audio_count == 0:
                audio_start = time.time()
            elif audio_count % 10 == 0:
                print "%dth audio: avg %.3g packets/sec." % (audio_count, (audio_count / (time.time() - audio_start)))
            audio_count += 1'''

        # Drop 1/30 audio packets
        '''if mtype == MTYPE_AUDIO and not random.randint(0, 30):
            print "Discarding audio packet! (%d bytes)" % buf.pos
            buf.discard()'''
        # Drop 1/30 video packets
        '''if mtype == MTYPE_VIDEO and not random.randint(0, 30):
            print "Discarding video packet! (%d bytes)" % buf.pos
            buf.discard()'''
        # Duplicate 1/30 Video packets
        '''if mtype == MTYPE_VIDEO and not random.randint(0, 30):
            print "Duplicating video packet x%d! (%d bytes)" % (dup_count, buf.pos)
            pos = buf.pos
            buf.seek(0)
            pack = buf.read(pos)
            buf.insert(pack * dup_count)'''
        # Duplicate 1/30 Video packets
        '''if mtype == MTYPE_AUDIO and not random.randint(0, 30):
            print "Duplicating audio packet x%d! (%d bytes)" % (dup_count, buf.pos)
            pos = buf.pos
            buf.seek(0)
            pack = buf.read(pos)
            buf.insert(pack * dup_count)
            dup_count += 1'''
        
        # Drop several audio/video packets in a row
        '''if i > 600 and dropped_packets < 500 and (mtype in (MTYPE_VIDEO, MTYPE_AUDIO)):
            print "Dropping %s packet! (%d bytes)" % (mtype == MTYPE_AUDIO and "audio" or "video", buf.pos)
            buf.discard()
            dropped_packets += 1'''

        # Duplicate a several audio/video packets and insert later on
        '''dup_max = 150
        if i > 600 and dup_count < dup_max and (mtype in (MTYPE_VIDEO, MTYPE_AUDIO)):
            print "Storing a/v packet for later! (%d bytes)" % buf.pos, i
            pos = buf.pos
            buf.seek(0)
            dupdata += buf.read(pos)
            dup_count += 1
        if i == 1000:
            print "Inserting duplicate a/v data into stream"
            buf.insert(dupdata)'''
        buf.push()

    #print "Forwarding data without analysis..."
    while 1:
        buf.read(1)
        buf.seek(-1)
        #print "srv -> cli: %d bytes" % buf.pos
        buf.push()

server_connection_dropped = False
streaming_state_received = False


def interceptClientBound(buf):
    '''
    Simulate failover: drop client bound traffic after 10 seconds,
    send get_state() to client, run StreamingFailover to open a
    new server connection and continue filling client's stream
    '''
    # forward handshake packets
    buf.pushn(1)
    buf.pushn(1536)
    buf.pushn(1536)

    """rtmp = RTMPStream(buf)

    loading = False
    comparing = False

    class channel:
        def __init__(self, name, dbfile):
           self.name = name
           self.map = PersistentDict(dbfile)
           self.time = 0
           self.seq = 0
           #self.started_failover = False
           #self.finished_failover = False
           #self.counter = 100
           self.prev_packet = ''
           #self.print_modified_packet = False

    audiochannel = channel('AUDIO', 'audiomap2.txt')
    videochannel = channel('VIDEO', 'videomap2.txt')
    channels = (audiochannel, videochannel)

    if loading:
        for channel in channels:
            channel.map.clear()

    #print "Forwarding data without analysis..."
    try:
        for i in range(1200):
            packet = rtmp.read_chunk()
            header = packet.header
            #log.debug(header)
            #packet_len = buf.pos
            #buf.seek(0)
            #packet_raw = buf.read(packet_len)
            #log.debug('raw = %s\n%s', binascii.hexlify(packet_raw), stringDump(packet_raw))
            #buf.push()
            #continue
        
            if header.mtype in (MTYPE_AUDIO, MTYPE_VIDEO) and i > -1:
                channel = {MTYPE_AUDIO : audiochannel,
                           MTYPE_VIDEO : videochannel}.get(header.mtype)

                if not comparing:
                    channel.seq += 1
                    log.debug('%s %d msec (delta: %d)', channel.name, header.time, header.time - channel.time) # stringDump(packet[:40])
                    channel.time = header.time
                if loading:
                    timeseq = channel.map.get(packet.data, [])
                    timeseq.append((header.time, channel.seq))
                    channel.map[packet.data] = timeseq
                elif comparing:
                    entries = channel.map.get(packet.data, [])
                    timeseq = ['%d msec (#%d)' % (e[0], e[1]) for e in entries[:4]]
                    log.debug('%d msec --> %s --> %s', header.time, channel.name, '; '.join(timeseq))
                    
                    #channel.counter += 1

                    '''if channel.finished_failover:
                        # correct packet time
                        fmt = ord(packet[0]) >> 6
                        if fmt == 0 and entries:
                            buf.seek(1)
                            buf.write(str24BE(entries[0][0]))
                            buf.seek(0)
                            packet = rtmp.decode()
                            packet_len = buf.pos
                            buf.seek(0)
                            packet_raw = buf.read(packet_len)
                            #if print_modified_packet:
                            #    buf.seek(0)
                            #    packet = buf.read(packet_len)
                            #    print '\n', stringDump(packet_raw)
                            #    channel.print_modified_packet = False'''

                    '''if not channel.finished_failover and\
                       (channel.started_failover and len(entries) == 1 and entries[0][1] > channel.seq):
                        channel.finished_failover = True
                        print stringDump(packet_raw)
                        channel.print_modified_packet = True
                        channel.counter = 0
                        
                    if not channel.started_failover and\
                       (i > 100 and (len(entries) == 0 or (len(entries) == 1 and entries[0][1] < (channel.seq - 1)))):
                        channel.started_failover = True
                        print stringDump(channel.prev_packet)
                        channel.counter = 0

                    if not channel.started_failover or channel.finished_failover:
                        channel.seq = max([e[1] for e in entries] + [0])'''

                    '''to_be_dropped = (channel.started_failover and not channel.finished_failover)\
                                    or (channel.counter < 10 and len(entries) == 0)
                    timeseq = ['%d msec (#%d)' % (e[0], e[1]) for e in entries[:5]]
                    #timeseq = '?'
                    if header.mtype == MTYPE_VIDEO:
                        log.debug('%d msec --> %s --> %s%s', header.time, channel.name, '; '.join(timeseq), to_be_dropped and ' ...DROPPED!' or '')
                    # "fix" the time to the expected time (stored in database)
                    if to_be_dropped:
                        buf.discard()'''

                    channel.prev_packet = packet


                    '''else:
                        fmt = ord(packet[0]) >> 6
                        if fmt == 0:
                            buf.seek(1)
                            buf.write(str24BE(times[0]))
                            #packet = packet[0] + str24BE(times[0]) + packet[4:]
                            buf.seek(0)
                            continue'''

            buf.push()
    finally:
        if loading:
            for channel in channels:
                channel.map.save()
                #for k, v in channel.map.iteritems():
                #    print '%s\n---> %s\n' % (stringDump(k), v)"""

    while 1:
        buf.read(1)
        buf.seek(-1)
        buf.push()

    # Inject RPC get_state()
    print "Injecting RPC message 'get_state()'"
    cmd = "03" + "000000" + "000017" + "11" + "00000000" + "000200096765745f737461746500000000000000000005"
    buf.seek(0)
    buf.insert(binascii.unhexlify(cmd))
    buf.push()

    # Decode next 3 packets
    '''for i in range(3):
        mtype, sid, payload = rtmp.decode()
        if mtype in (MTYPE_AMF0_COMMAND, MTYPE_AMF3_COMMAND):
            
        buf.push()'''

def interceptServerBound(buf):
    '''Watches data from client to server and inserts or removes data as desired.'''

    # forward handshake packets
    buf.pushn(1)
    buf.pushn(1536)
    buf.pushn(1536)

    rtmp = RTMPStream(buf)

    decode_enabled = False

    if not decode_enabled:
        while 1:
            buf.read(1)
            buf.seek(-1)
            buf.push()
    else:
        while 1:
            packet = rtmp.read_packet()
            buf.push()

    '''# Decode next 3 packets
    for i in xrange(100000):
        t, mtype, sid, payload = create_streamdecode()
        # expecting receive_state()
        if server_connection_dropped and mtype == MTYPE_AMF3_COMMAND:
            print '\n', MTYPES[mtype], sid
            print stringDump(payload)
            pos = buf.pos
            buf.seek(0)
            streaming_state_packet = buf.read(pos)
            streaming_state = fo_client.decode_state(streaming_state_packet)
            break
        else:
            buf.push()'''

    '''f = fo_client.StreamingFailover(streaming_state, random.randint(50000, 60000), 'Failover #M')
    f.failover()
    new_srv_s = f.sock

    # forward data between client and server using old client socket and new
    # server socket forever...
    m.deactivate()
    new_srv_s.setblocking(0)
    cli_s.setblocking(0)
    while 1:
        new_srv_s.sendall(cli_s.recv(4096))
        cli_s.sendall(new_srv_s.recv(4096))'''

    #print "Forwarding data without analysis..."
    '''while 1:
        buf.read(1)
        buf.seek(-1)
        buf.push()'''

if __name__ == '__main__':
    m = MiddleMan(interceptClientBound, interceptServerBound)
    m.mediate(2011, 1967, 'localhost')
    assert False, "MiddleMan should never finish!"
